{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checkpointing Workflow Runs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrate how to checkpoint `Workflow` runs via a `WorkflowCheckpointer` object. We also show how we can view all of the checkpoints that are stored in this object and finally how we can use a checkpoint as the starting point of a new run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "api_key = os.environ.get(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Workflow,\n",
    "    step,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Event,\n",
    "    Context,\n",
    ")\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "\n",
    "class JokeEvent(Event):\n",
    "    joke: str\n",
    "\n",
    "\n",
    "class JokeFlow(Workflow):\n",
    "    llm = OpenAI(api_key=api_key)\n",
    "\n",
    "    @step\n",
    "    async def generate_joke(self, ev: StartEvent) -> JokeEvent:\n",
    "        topic = ev.topic\n",
    "\n",
    "        prompt = f\"Write your best joke about {topic}.\"\n",
    "        response = await self.llm.acomplete(prompt)\n",
    "        return JokeEvent(joke=str(response))\n",
    "\n",
    "    @step\n",
    "    async def critique_joke(self, ev: JokeEvent) -> StopEvent:\n",
    "        joke = ev.joke\n",
    "\n",
    "        prompt = f\"Give a thorough analysis and critique of the following joke: {joke}\"\n",
    "        response = await self.llm.acomplete(prompt)\n",
    "        return StopEvent(result=str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a WorkflowCheckpointer Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow.checkpointer import WorkflowCheckpointer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate Jokeflow\n",
    "workflow = JokeFlow()\n",
    "wflow_ckptr = WorkflowCheckpointer(workflow=workflow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Workflow from the WorkflowCheckpointer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `WorkflowCheckpointer.run()` method is a wrapper over the `Workflow.run()` method, which injects a checkpointer callback in order to create and store checkpoints. Note that checkpoints are created at the completion of a step, and that the data stored in checkpoints are:\n",
    "\n",
    "- `last_completed_step`: The name of the last completed step\n",
    "- `input_event`: The input event to this last completed step\n",
    "- `output_event`: The event outputted by this last completed step\n",
    "- `ctx_state`: a snapshot of the attached `Context`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This joke plays on the double meaning of the word \"rates,\" which can refer to both the cost of something and the passage of time. The punchline suggests that chemists prefer nitrates because they are less expensive than day rates, implying that chemists are frugal or cost-conscious individuals.\\n\\nOverall, the joke is clever and plays on a pun that is likely to be appreciated by those familiar with chemistry and the concept of nitrates. However, the humor may be lost on those who are not well-versed in chemistry terminology. Additionally, the joke relies on a somewhat simplistic play on words, which may not be as engaging or humorous to some audiences.\\n\\nIn terms of structure, the joke follows a classic setup and punchline format, with the punchline providing a surprising twist on the initial premise. The delivery of the joke may also play a role in its effectiveness, as timing and tone can greatly impact the humor of a joke.\\n\\nOverall, while the joke may appeal to a specific audience and demonstrate some clever wordplay, it may not have universal appeal and may be considered somewhat niche in its humor.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "handler = wflow_ckptr.run(\n",
    "    topic=\"chemistry\",\n",
    ")\n",
    "await handler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can view all of the checkpoints via the `.checkpoints` attribute, which is dictionary with keys representing the `run_id` of the run and whose values are the list of checkpoints stored for the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'483eccdd-a035-42cc-b596-cd33d42938a7': [Checkpoint(id_='d5acd098-47e2-4acf-9520-9ca06ee4e238', last_completed_step='generate_joke', input_event=StartEvent(), output_event=JokeEvent(joke=\"Why do chemists like nitrates so much?\\n\\nBecause they're cheaper than day rates!\"), ctx_state={'globals': {}, 'streaming_queue': '[]', 'queues': {'_done': '[]', 'critique_joke': '[]', 'generate_joke': '[]'}, 'stepwise': False, 'events_buffer': {}, 'in_progress': {'generate_joke': []}, 'accepted_events': [('generate_joke', 'StartEvent'), ('critique_joke', 'JokeEvent')], 'broker_log': ['{\"__is_pydantic\": true, \"value\": {\"_data\": {\"topic\": \"chemistry\", \"store_checkpoints\": false}}, \"qualified_name\": \"llama_index.core.workflow.events.StartEvent\"}'], 'is_running': True}),\n",
       "  Checkpoint(id_='288c3e54-292b-4c7e-aed8-662537508b46', last_completed_step='critique_joke', input_event=JokeEvent(joke=\"Why do chemists like nitrates so much?\\n\\nBecause they're cheaper than day rates!\"), output_event=StopEvent(result='This joke plays on the double meaning of the word \"rates,\" which can refer to both the cost of something and the passage of time. The punchline suggests that chemists prefer nitrates because they are less expensive than day rates, implying that chemists are frugal or cost-conscious individuals.\\n\\nOverall, the joke is clever and plays on a pun that is likely to be appreciated by those familiar with chemistry and the concept of nitrates. However, the humor may be lost on those who are not well-versed in chemistry terminology. Additionally, the joke relies on a somewhat simplistic play on words, which may not be as engaging or humorous to some audiences.\\n\\nIn terms of structure, the joke follows a classic setup and punchline format, with the punchline providing a surprising twist on the initial premise. The delivery of the joke may also play a role in its effectiveness, as timing and tone can greatly impact the humor of a joke.\\n\\nOverall, while the joke may appeal to a specific audience and demonstrate some clever wordplay, it may not have universal appeal and may be considered somewhat niche in its humor.'), ctx_state={'globals': {}, 'streaming_queue': '[]', 'queues': {'_done': '[]', 'critique_joke': '[]', 'generate_joke': '[]'}, 'stepwise': False, 'events_buffer': {}, 'in_progress': {'generate_joke': [], 'critique_joke': []}, 'accepted_events': [('generate_joke', 'StartEvent'), ('critique_joke', 'JokeEvent')], 'broker_log': ['{\"__is_pydantic\": true, \"value\": {\"_data\": {\"topic\": \"chemistry\", \"store_checkpoints\": false}}, \"qualified_name\": \"llama_index.core.workflow.events.StartEvent\"}', '{\"__is_pydantic\": true, \"value\": {\"joke\": \"Why do chemists like nitrates so much?\\\\n\\\\nBecause they\\'re cheaper than day rates!\"}, \"qualified_name\": \"__main__.JokeEvent\"}'], 'is_running': True})]}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wflow_ckptr.checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run: 483eccdd-a035-42cc-b596-cd33d42938a7 has 2 stored checkpoints\n"
     ]
    }
   ],
   "source": [
    "for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
    "    print(f\"Run: {run_id} has {len(ckpts)} stored checkpoints\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filtering the Checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `WorkflowCheckpointer` object also has a `.filter_checkpoints()` method that allows us to filter via:\n",
    "\n",
    "- The name of the last completed step by speciying the param `last_completed_step`\n",
    "- The event type of the last completed step's output event by specifying `output_event_type`\n",
    "- Similarly, the event type of the last completed step's input event by specifying `input_event_type`\n",
    "\n",
    "Specifying multiple of these filters will be combined by the \"AND\" operator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test this functionality out, but first we'll make things a bit more interesting by running a couple of more runs with our `Workflow`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "additional_topics = [\"biology\", \"history\"]\n",
    "\n",
    "for topic in additional_topics:\n",
    "    handler = wflow_ckptr.run(topic=topic)\n",
    "    await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run: 483eccdd-a035-42cc-b596-cd33d42938a7 has 2 stored checkpoints\n",
      "Run: e112bca9-637c-4492-a8aa-926c302c99d4 has 2 stored checkpoints\n",
      "Run: 7a59c918-90a3-47f8-a818-71e45897ae39 has 2 stored checkpoints\n"
     ]
    }
   ],
   "source": [
    "for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
    "    print(f\"Run: {run_id} has {len(ckpts)} stored checkpoints\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Checkpoint(id_='d5acd098-47e2-4acf-9520-9ca06ee4e238', last_completed_step='generate_joke', input_event=StartEvent(), output_event=JokeEvent(joke=\"Why do chemists like nitrates so much?\\n\\nBecause they're cheaper than day rates!\"), ctx_state={'globals': {}, 'streaming_queue': '[]', 'queues': {'_done': '[]', 'critique_joke': '[]', 'generate_joke': '[]'}, 'stepwise': False, 'events_buffer': {}, 'in_progress': {'generate_joke': []}, 'accepted_events': [('generate_joke', 'StartEvent'), ('critique_joke', 'JokeEvent')], 'broker_log': ['{\"__is_pydantic\": true, \"value\": {\"_data\": {\"topic\": \"chemistry\", \"store_checkpoints\": false}}, \"qualified_name\": \"llama_index.core.workflow.events.StartEvent\"}'], 'is_running': True}),\n",
       " Checkpoint(id_='87865641-14e7-4eb0-bb62-4c211567acfc', last_completed_step='generate_joke', input_event=StartEvent(), output_event=JokeEvent(joke=\"Why did the biologist break up with the mathematician?\\n\\nBecause they couldn't find a common denominator!\"), ctx_state={'globals': {}, 'streaming_queue': '[]', 'queues': {'_done': '[]', 'critique_joke': '[]', 'generate_joke': '[]'}, 'stepwise': False, 'events_buffer': {}, 'in_progress': {'generate_joke': []}, 'accepted_events': [('generate_joke', 'StartEvent'), ('critique_joke', 'JokeEvent')], 'broker_log': ['{\"__is_pydantic\": true, \"value\": {\"_data\": {\"topic\": \"biology\"}}, \"qualified_name\": \"llama_index.core.workflow.events.StartEvent\"}'], 'is_running': True}),\n",
       " Checkpoint(id_='69a99535-d45c-46b4-a1f4-d3ecc128cb08', last_completed_step='generate_joke', input_event=StartEvent(), output_event=JokeEvent(joke='Why did the history teacher go to the beach?\\n\\nTo catch some waves of the past!'), ctx_state={'globals': {}, 'streaming_queue': '[]', 'queues': {'_done': '[]', 'critique_joke': '[]', 'generate_joke': '[]'}, 'stepwise': False, 'events_buffer': {}, 'in_progress': {'generate_joke': []}, 'accepted_events': [('generate_joke', 'StartEvent'), ('critique_joke', 'JokeEvent')], 'broker_log': ['{\"__is_pydantic\": true, \"value\": {\"_data\": {\"topic\": \"history\"}}, \"qualified_name\": \"llama_index.core.workflow.events.StartEvent\"}'], 'is_running': True})]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Filter by the name of last completed step\n",
    "checkpoints_right_after_generate_joke_step = wflow_ckptr.filter_checkpoints(\n",
    "    last_completed_step=\"generate_joke\",\n",
    ")\n",
    "\n",
    "# checkpoint ids\n",
    "[ckpt for ckpt in checkpoints_right_after_generate_joke_step]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Re-Run Workflow from a specific checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run from a chosen `Checkpoint` we can use the `WorkflowCheckpointer.run_from()` method. NOTE that doing so will lead to a new `run` and it's checkpoints if enabled will be stored under the newly assigned `run_id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Analysis:\\nThis joke plays on the double meaning of the word \"rates,\" which can refer to both the cost of something and the passage of time. In this case, the joke suggests that chemists prefer nitrates because they are less expensive than day rates, implying that chemists are frugal or cost-conscious individuals.\\n\\nCritique:\\n- Clever wordplay: The joke relies on a clever play on words, which can be entertaining for those who appreciate puns and linguistic humor.\\n- Niche audience: The humor in this joke may be more appreciated by individuals with a background in chemistry or a specific interest in science, as the punchline relies on knowledge of chemical compounds.\\n- Lack of universal appeal: The joke may not resonate with a general audience who may not understand the reference to nitrates or the significance of their cost compared to day rates.\\n- Lack of depth: While the joke is amusing on a surface level, it may be considered somewhat shallow or simplistic compared to more nuanced or thought-provoking humor.\\n\\nOverall, the joke is a light-hearted play on words that may appeal to individuals with a specific interest in chemistry or wordplay. However, its niche appeal and lack of universal relevance may limit its effectiveness as a joke for a broader audience.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# can work with a new instance\n",
    "new_workflow_instance = JokeFlow()\n",
    "wflow_ckptr.workflow = new_workflow_instance\n",
    "\n",
    "ckpt = checkpoints_right_after_generate_joke_step[0]\n",
    "\n",
    "handler = wflow_ckptr.run_from(checkpoint=ckpt)\n",
    "await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run: 483eccdd-a035-42cc-b596-cd33d42938a7 has 2 stored checkpoints\n",
      "Run: e112bca9-637c-4492-a8aa-926c302c99d4 has 2 stored checkpoints\n",
      "Run: 7a59c918-90a3-47f8-a818-71e45897ae39 has 2 stored checkpoints\n",
      "Run: 9dccfda3-b5bf-4771-8293-efd3e3a275a6 has 1 stored checkpoints\n"
     ]
    }
   ],
   "source": [
    "for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
    "    print(f\"Run: {run_id} has {len(ckpts)} stored checkpoints\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we've executed from the checkpoint that represents the end of \"generate_joke\" step, there is only one additional checkpoint (i.e., that for the completion of step \"critique_joke\") that gets stored in the last partial run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying Which Steps To Checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default all steps of the attached workflow (excluding the \"_done\" step) will be checkpointed. You can see which steps are enabled for checkpointing via the `enabled_checkpoints` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'critique_joke', 'generate_joke'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wflow_ckptr.enabled_checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To disable a step for checkpointing, we can use the `.disable_checkpoint()` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wflow_ckptr.disable_checkpoint(step=\"critique_joke\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Analysis:\\nThis joke plays on the common stereotype that mechanics are overly critical and nitpicky when it comes to fixing cars. The humor comes from the unexpected twist of the car being the one to break up with the mechanic, rather than the other way around. The punchline is clever and plays on the idea of a relationship ending due to constant criticism.\\n\\nCritique:\\nOverall, this joke is light-hearted and easy to understand. It relies on a simple pun and doesn't require much thought to appreciate. However, the humor may be seen as somewhat predictable and not particularly original. The joke also perpetuates the stereotype of mechanics being overly critical, which may not sit well with some people in the automotive industry. Additionally, the joke may not be as universally funny as some other jokes, as it relies on a specific understanding of the relationship between cars and mechanics. Overall, while the joke is amusing, it may not be considered particularly groundbreaking or memorable.\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "handler = wflow_ckptr.run(topic=\"cars\")\n",
    "await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run: 483eccdd-a035-42cc-b596-cd33d42938a7 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: e112bca9-637c-4492-a8aa-926c302c99d4 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: 7a59c918-90a3-47f8-a818-71e45897ae39 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: 9dccfda3-b5bf-4771-8293-efd3e3a275a6 has stored checkpoints for steps ['critique_joke']\n",
      "Run: 6390e2ce-63f4-44a3-8a75-64ccb765abfd has stored checkpoints for steps ['generate_joke']\n"
     ]
    }
   ],
   "source": [
    "for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
    "    print(\n",
    "        f\"Run: {run_id} has stored checkpoints for steps {[c.last_completed_step for c in ckpts]}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can turn checkpointing back on by using the `.enable_checkpoint()` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wflow_ckptr.enable_checkpoint(step=\"critique_joke\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Analysis:\\nThis joke plays on the common stereotype that mechanics are overly critical and nitpicky when it comes to fixing cars. The humor comes from the unexpected twist of the car being the one to break up with the mechanic, rather than the other way around. The joke also cleverly uses the term \"nitpicking\" in a literal sense, as in picking at the nit (small details) of the car.\\n\\nCritique:\\nWhile the joke is clever and plays on a well-known stereotype, it may not be the most original or groundbreaking joke. The punchline is somewhat predictable and relies on a common trope about mechanics. Additionally, the joke may not be universally relatable or understood by all audiences, as it requires some knowledge of the stereotype about mechanics being nitpicky.\\n\\nOverall, the joke is light-hearted and humorous, but it may not be the most memorable or impactful joke due to its reliance on a common stereotype. It could be improved by adding a more unexpected or unique twist to the punchline.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "handler = wflow_ckptr.run(topic=\"cars\")\n",
    "await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run: 483eccdd-a035-42cc-b596-cd33d42938a7 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: e112bca9-637c-4492-a8aa-926c302c99d4 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: 7a59c918-90a3-47f8-a818-71e45897ae39 has stored checkpoints for steps ['generate_joke', 'critique_joke']\n",
      "Run: 9dccfda3-b5bf-4771-8293-efd3e3a275a6 has stored checkpoints for steps ['critique_joke']\n",
      "Run: 6390e2ce-63f4-44a3-8a75-64ccb765abfd has stored checkpoints for steps ['generate_joke']\n",
      "Run: e3291623-6eb8-43c6-b102-dc6f88a42f4d has stored checkpoints for steps ['generate_joke', 'critique_joke']\n"
     ]
    }
   ],
   "source": [
    "for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
    "    print(\n",
    "        f\"Run: {run_id} has stored checkpoints for steps {[c.last_completed_step for c in ckpts]}\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-core-nrCJU7vQ-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
